import os
from typing import Optional

from models.clip_model import ClipModel
from models.knowledge_base_lite import KnowledgeBaseLite

# 检查是否安装了完整版pymilvus
try:
    from models.knowledge_base import KnowledgeBase
    HAS_FULL_MILVUS = True
except ImportError:
    HAS_FULL_MILVUS = False

class KnowledgeBaseFactory:
    """知识库工厂类，用于创建适当的知识库实现"""
    
    @staticmethod
    def create_knowledge_base(
        base_path: str = "./knowledge_base",
        clip_model: Optional[ClipModel] = None,
        use_lite: bool = False,
        milvus_host: str = "localhost",
        milvus_port: str = "19530"
    ):
        """
        创建知识库实例
        
        Args:
            base_path: 知识库存储路径
            clip_model: CLIP模型实例
            use_lite: 是否使用轻量版实现
            milvus_host: Milvus服务器地址
            milvus_port: Milvus服务器端口
            
        Returns:
            知识库实例
        """
        # 如果指定使用轻量版或者没有安装完整版pymilvus，使用轻量版实现
        if use_lite or not HAS_FULL_MILVUS or milvus_host == "lite":
            print("使用Milvus Lite本地向量存储")
            return KnowledgeBaseLite(base_path=base_path, clip_model=clip_model)
        
        # 否则使用完整版实现
        print(f"使用Milvus服务器向量存储: {milvus_host}:{milvus_port}")
        return KnowledgeBase(
            base_path=base_path,
            clip_model=clip_model,
            milvus_host=milvus_host,
            milvus_port=milvus_port
        ) 